from typing import List, Union
import numpy as np
from fixedint import UInt64

def join_words(words: List[int]) -> int:
    """Join a list of 64bit words into a single python integer"""
    out_str = ''
    for x in words[::-1]:
        out_str += f"{x:064b}"
    return int(out_str, 2)

def split_into_words(x: str) -> List[int]:
    """Split a binary representation of a python int into a list of 64 bit words"""
    strings = [x[i:i+64] for i in range(0, len(x), 64)]
    return [int(i, 2) for i in strings[::-1]]

def adder(lhs: List[UInt64], rhs: List[UInt64]) -> List[UInt64]:
    """Multi-word adder"""
    carry = 0
    for i, _ in enumerate(lhs):
        vbf = lhs[i]
        lhs[i] = lhs[i] + rhs[i] + carry
        carry = lhs[i] < vbf or (carry & (lhs[i] == vbf))
    return lhs


def to_fixed(x: int, width=192):
    """Fixed width string rep of python integer"""
    return f"{x:0{width}b}"[-width:]

def to_cpp_vec(x: List[Union[int, UInt64]], label="x"):
    """Write as a vector of integers"""
    out_str = f"std::vector<uint64_t> {label}{{"
    for i, w in enumerate(x):
        comma = "," if i < len(x) - 1 else ""
        out_str += f"{w}ULL{comma}"
    out_str += "};"
    print(out_str)

def gen_random_integer(num_bits):
    num_words = (num_bits - 1) // 64 + 1
    words = []
    for iw in range(num_words):
        r = np.random.random()
        b = 0
        if r < 1/3:
            b = 0
        elif r > 2/3:
            if iw == num_words - 1:
                for k in range(num_bits - 64 * iw):
                    b |= (1<< k)
            else:
                b = 0xFFFFFFFFFFFFFFFF
        else:
            for k in range(64):
                if np.random.random() > 0.5:
                    b |= (1<< k)
        words.append(b)
    return words

if __name__ == "__main__":
    np.random.seed(7)
    #for i in range(10_000):
    for i in range(4):
        #num_bits = np.random.randint(10, 1200) 
        num_bits = 512
        x = gen_random_integer(num_bits)
        y = gen_random_integer(num_bits)
        a = join_words(x)
        b = join_words(y)
        xpy = to_fixed(a+b, num_bits)
        xu = [UInt64(i) for i in x]
        yu = [UInt64(i) for i in y]
        z = adder(xu, yu)
        z_str = to_fixed(join_words(z), num_bits)
        xpy_str = xpy
        if i == 3:
            to_cpp_vec(x, "x")
            to_cpp_vec(y, "y")
            to_cpp_vec(split_into_words(xpy), "z")
        assert to_fixed(join_words(z), num_bits) == xpy
        assert xpy_str == z_str